# coding=utf-8
import codecs
import datetime
import os
import shutil
import threading
from queue import Queue

import pandas
import pandas as pd
from loguru import logger

import main_cal_down
from cal_ops.point import cal_point1_txt, cal_point1_aa
from const import stocks
from models.stock_model import StockNumber
from mylib import download_all
from mylib.download_all import analysis_stock
from mylib.mycsv import sort_csv
from mylib.mydb import MyDB

st_queue = Queue()


def get_all_txt_files():
    for root, dirs, files in os.walk('txts'):
        return [(os.path.join(root, item), item.replace('.txt', '')) for item in files]


def clear_dir(d='result'):
    if os.path.exists(d):
        shutil.rmtree(d)
    os.makedirs(d)


def get_all_stock_csv_path():
    for root, dirs, files in os.walk('stocks'):
        return [os.path.join(root, item) for item in files]


def get_all_stocks_indus():
    with open('industry_count_rank.csv', 'r') as fr:
        all_indus = [item.strip().split(',') for item in fr.readlines()[1:]]
        return all_indus


def get_avg_top10(stocks, N=50):
    stocks_top_dict = {}
    for stock in stocks:
        csv_path = f'stocks/{stock}.csv'
        if not os.path.exists(csv_path):
            continue
        df2 = pd.read_csv(csv_path)
        stocks_top_dict[stock] = round(abs(df2[0:N]['pct_chg'].mean()), 4)
    sorted_d = sorted(stocks_top_dict.items(), key=lambda x: x[1], reverse=True)
    # 获取近N日振幅排名前10的股票
    return sorted_d[:10]


def get_end_date(log_dir, indus, indus_num):
    f_path = f'{log_dir}/{indus_num}_all_{indus}.csv'
    if not os.path.exists(f_path):
        return None
    with open(f_path, 'r') as fr:
        line = fr.readline()
        return eval(line.split(',')[0])


def t_download(N):
    try:
        df = pd.read_csv('all.csv')
        for row in df.index:
            if row > N:
                break
            if stocks and df.loc[row]['ts_code'] not in stocks:
                continue
            sn = StockNumber(df.loc[row])
            if '退' in sn.name:
                continue
            if 'ST' in sn.name:
                continue
            if not str(sn.ts_code).startswith('60') \
                    and not str(sn.ts_code).startswith('00'):
                continue
            analysis_stock(sn)
            st_queue.put(('cal_aa', 1, [sn.ts_code]))
    except Exception as e:
        logger.error(e)
    st_queue.put((1, 1, 1))


def t_cal(full_path_csv):
    log_dir = os.path.basename(__file__).split('.')[0]
    current_date = str(datetime.datetime.now().date()).replace('-', '')
    log_dir = f'{log_dir}/{current_date}'
    # 需要个股涨跌图打开，一般只计算一天才打开
    gen_excel_flag = True
    bkd = 1
    start_date = 20991230
    with open(full_path_csv, 'w', encoding='utf-8') as aa:
        aa.write(f'AA,BB,CC,DD,EE,FF,GG,HH,II,JJ')
        while 1:
            if not st_queue:
                continue
            indus, indus_num, stocks = st_queue.get()
            if (1, 1, 1) == (indus, indus_num, stocks):
                logger.warning('所有计算完成')
                break
            end_date = get_end_date(log_dir, indus, indus_num)
            logger.info(f'开始计算{stocks}')
            msg = cal_point1_aa.run(log_dir, start_date, end_date, stocks=stocks, bkd=bkd)
            aa.write(msg)
        logger.success('计算完成')


if __name__ == '__main__':
    logger.add("log/main_cal_today_{time}.log", level='WARNING')
    full_path_csv = 'aa.csv'
    N = 100
    t1 = threading.Thread(target=t_download, args=(N,))
    t2 = threading.Thread(target=t_cal, args=(full_path_csv,))
    t2.start()
    t1.start()
    t2.join()
    t1.join()
    sort_csv(full_path_csv, ['FF', 'HH', 'JJ'], [True, False, False])
    logger.success('排序完成')
    df = pandas.read_csv(full_path_csv)
    ss = [item for item in df['BB'][:20]]
    main_cal_down.run(ss, d=False)
    logger.success('main_cal_down计算完成')
